from typing import Optional, List, Mapping, Any
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_community.llms.base import LLM

import requests
import os

class BaiLianLLM(LLM):
    model_name: str = "qwen-max"
    api_key: str = os.getenv("DASHSCOPE_API_KEY")

    @property
    def _llm_type(self) -> str:
        return "bailian"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        data = {
            "model": self.model_name,
            "input": {
                "prompt": prompt
            }
        }

        response = requests.post(
            "https://api.bailian.aliyun.com/api/v1/services/aigc/text-generation/generation ",
            headers=headers,
            json=data
        )

        if response.status_code != 200:
            raise Exception(f"API Error: {response.text}")

        return response.json()["output"]["text"]

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        return {"model_name": self.model_name}